import requests
import datetime
import base64

host_port = 'http://8.130.16.54:17860'

model_name_list = ['meinahentai_v21.safetensors [deb40a068a]',
                   'ligneClaireAnime_v1.safetensors [672977e447]',
                   'cetusMix_v4.safetensors [b42b09ff12]']

size_list = [512, 768]
max_gen_num = 3

text2img_payload = {
    "prompt": "",
    "batch_size": 1,
    "negative_prompt": "easynegative, nsfw, nude",
    "sampler_index": "Euler a",
    "seed": -1,
    "steps": 25,
    "width": 512,
    "height": 512,
    "cfg_scale": 7,
}


def gen_file_name(file_type: str, plus_text: str, file_idx: int):
    date_str = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    return f'{plus_text}_{date_str}_{file_idx}{file_type}'


def get_sd_model_list() -> list[str]:
    url = f'{host_port}/sdapi/v1/sd-models'
    res = requests.get(url)
    if res.status_code != 200:
        raise Exception()
    return [model_info['title'] for model_info in res.json()]


def set_model(model_idx: int):
    r"""设置模型，三种模型可选。

    :param model_idx: 模型序号，列表如下：【0】meinahentai，【1】ligneClaireAnime，【2】cetusMix，请自己感受差别。
    :return: 无返回
    """
    if model_idx not in range(len(model_name_list)):
        print('[!] 输入非法！')
        raise ValueError('输入非法！')
    print(f'[*] 修改模型：{model_name_list[model_idx]}')
    url = f'{host_port}/sdapi/v1/options'
    payload = {'CLIP_stop_at_last_layers': 2, 'sd_model_checkpoint': model_name_list[model_idx]}
    res = requests.post(url, json=payload)
    if res.status_code != 200:
        print('[!]', res.json()['detail']['msg'])
    print('[*] 模型修改成功')


def text2img_base64(prompt: str, gen_num: int = 1, gen_width: int = 512, gen_height: int = 512) -> list[str]:
    r"""文字生成图片（文生图）。

    :param gen_height: 生成图片的高（像素），取值：512，768
    :param gen_width: 生成图片的宽（像素），取值：512，768
    :param gen_num: 生成图片个数，取值：1-3
    :param prompt: 描述图片的一个字符串，形式多样，可以是一个句子：“a girl with red eyes and white hair with cat ear”，
                也可以是所有关键词用逗号隔开：”1girl, red eyes, white hair, cat ear“。
    :return: 一个列表，列表的每一项为生成图片二进制数据的的base64编码字符串。
    """
    print('[*] 加载参数列表')
    url = f'{host_port}/sdapi/v1/txt2img'
    payload = text2img_payload
    payload['prompt'] = prompt
    if gen_num not in range(1, max_gen_num + 1):
        print(f'[!] 生成数量应在 1到{max_gen_num}之间!')
        raise ValueError("数值非法！")
    payload['batch_size'] = gen_num
    if gen_width not in size_list or gen_height not in size_list:
        print(f'[!] 图像的宽高只能取值为 {size_list}!')
        raise ValueError("数值非法！")
    payload['width'] = gen_width
    payload['height'] = gen_height
    print('[*] 生成中，请稍等')
    res = requests.post(url, json=payload)
    if res.status_code != 200:
        print('[!]', res.json()['detail']['msg'])
    print(res.json()['parameters'])
    print('[*] 已获取图像base64转码数据')
    return res.json()['images']


def text2img_save(prompt: str, gen_num: int = 1, gen_width: int = 512, gen_height: int = 512, save_dir: str = './'):
    r"""文字生成图片（文生图）,并保存至指定文件夹。

    :param gen_height: 生成图片的高（像素），取值：512，768
    :param gen_width: 生成图片的宽（像素），取值：512，768
    :param gen_num: 生成图片个数，取值：1-3
    :param prompt: 描述图片的一个字符串，形式多样，可以是一个句子：“a girl with red eyes and white hair with cat ear”，
                也可以是所有关键词用逗号隔开：”1girl, red eyes, white hair, cat ear“。
    :param save_dir: 保存文件夹路径，默认为同级目录内。
    :return: 无返回
    """
    base64_list = text2img_base64(prompt, gen_num, gen_width, gen_height)
    for idx, base64_img_str in enumerate(base64_list):
        img_data = base64.b64decode(base64_img_str)
        file_name = gen_file_name('.png', 'text2img', idx)
        print(f'[*] 文件保存为： {save_dir}{file_name}')
        with open(f'{save_dir}{file_name}', 'wb') as f:
            f.write(img_data)
    print('[*] 结束')


if __name__ == '__main__':
    pass
    # set_model(2)
    # text2img_save("1girl, solo, black hair, black eyes, smile, closed mouth, masterpiece, best quality"
    #               , gen_height=512, gen_width=768, gen_num=2)
